package rpcbiz

import (
	"errors"
	"github.com/golang/glog"
	"strings"
	"fmt"
	"time"
	"cloud_worm/util/unet/rpc"
	"cloud_worm/util/proto/urpc"
	"cloud_worm/util/uconst"
	"cloud_worm/access/server/client"
	"cloud_worm/util/proto/ujson"
	"cloud_worm/util/proto/uweapon"
)

// --------------------- cmd exec ---------------------
type CmdArgs urpc.CmdArgs
// CmdArgs.cmdParam 参数处理
func (args *CmdArgs) cmdParam() error {

	if len(args.Id) != uconst.AgentIdLen {
		glog.Errorf("input worm id error: %s", args.Id)
		return errors.New("input worm id error")
	}

	if len(args.TaskId) != uconst.TaskIdLen {
		glog.Errorf("input task id error: %s", args.TaskId)
		return errors.New("input task id error")
	}

	if len(args.Cmd) <= 0 {
		glog.Errorf("input cmd error: task_id=%s, wolid=%s", args.TaskId, args.Id)
		return errors.New("input cmd error")
	}

	if len(args.Params) > 0 {
		params := strings.Split(args.Params, " ")
		if len(params) > uconst.CmdParamMaxNum {
			glog.Error(fmt.Sprintf("'params' num > %d", uconst.CmdParamMaxNum))
			return fmt.Errorf("'params' num > %d", uconst.CmdParamMaxNum)
		}
	}

	if args.TimeOut <= uconst.WpTimeoutValid {
		args.TimeOut = uconst.WpTimeoutValid
	}
	return nil
}

// ExecAction.ExecCmd 命令执行
func (e *ExecAction) ExecCmd(args *CmdArgs, reply *urpc.CmdReply) error {
	// 参数解析
	err := args.cmdParam()
	if err != nil {
		return err
	}

	// 处理
	if client.AgentPoolInst.HaveConnByID(args.Id) == false {
		glog.Errorf("'%s' offline", args.Id)
		return errors.New("agent offline")
	}

	toAgentCmdMsg := ujson.MsgCmd{
		Id: args.Id,
		TaskId: args.TaskId,
		Cmd: args.Cmd,
		Params: args.Params,
		Timeout: args.TimeOut,
	}
	content, _ := toAgentCmdMsg.MarshalJSON()
	sendMsg := ujson.ProtoMsg{Type: ujson.ProtoCmdMsg, Content: content}
	buf, err := sendMsg.MarshalJSON()
	if err != nil {
		return err
	}

	// 发送send
	agent := client.AgentPoolInst.GetAgentByID(args.Id)
	err = agent.Connection.Write(buf)
	if err != nil {
		glog.Errorf("Send CmdMsg Error: %s set offline", agent.IP)
		client.AgentPoolInst.DelAgent(agent)
		return errors.New("send cmd msg error")
	}

	// 创建一个channel
	agent.CmdChanLock.Lock()
	ch := make(chan interface{})
	agent.CmdChanMap[args.TaskId] = ch       // task id是唯一性的
	agent.CmdChanLock.Unlock()
	// 等待结果
	err = nil
	select {
	case cmdRet := <- ch:
		reply.TaskId = cmdRet.(uweapon.CmdBrief).TaskId
		reply.Status = cmdRet.(uweapon.CmdBrief).Status
		reply.Result = cmdRet.(uweapon.CmdBrief).Result
		reply.Input = cmdRet.(uweapon.CmdBrief).Input
		break
	case <- time.After(time.Second * 60):
		// 60s秒超时，关闭channel
		glog.Errorf("'%s' exec '%s' timeout", agent.IP, args.Cmd)
		err = errors.New("exec time out")
		break
	}
	agent.CmdChanLock.Lock()
	delete(agent.CmdChanMap, args.TaskId)            // 删除map资源
	close(ch)                                               // 关闭channel
	agent.CmdChanLock.Unlock()

	return err
}

// --------------------- weapon exec ---------------------
// WeaponMsgArgs weapon执行相关
type WeaponMsgArgs urpc.WeaponMsgArgs
// WeaponMsgArgs.wpMsgParam 参数处理
func (args *WeaponMsgArgs) wpMsgParam() error {
	if len(args.Id) != uconst.AgentIdLen {
		glog.Errorf("input worm id error")
		return errors.New("input worm id error")
	}
	if len(args.TaskId) != uconst.TaskIdLen {
		glog.Errorf("'%s' input task id error", args.Id)
		return errors.New("input task id error")
	}
	if args.WeaponId <= 0 {
		return errors.New("input weapon id error")
	}
	if len(args.WeaponName) <= 0 {
		return errors.New("input weapon name error")
	}
	if len(args.WeaponMd5) != uconst.WpMd5Len {
		return errors.New("input weapon md5 error")
	}
	if !strings.HasPrefix(args.WeaponUrl, "https://") && !strings.HasPrefix(args.WeaponUrl, "http://") {
		glog.Errorf("'%s' input weapon url error: %s", args.Id, args.WeaponUrl)
		return errors.New("input weapon url error")
	}
	if args.Timeout <= uconst.WpTimeoutValid {
		args.Timeout = uconst.WpTimeoutValid
	}
	return nil
}

// ExecAction.ExecWp 流程处理
func (e *ExecAction) ExecWp(args *WeaponMsgArgs, reply *urpc.WeaponMsgReply) error {
	// 参数处理
	err := args.wpMsgParam()
	if err != nil {
		return err
	}

	// 流程处理
	if client.AgentPoolInst.HaveConnByID(args.Id) == false {
		return errors.New("agent offline")
	}
	// 检查在线，进行业务处理
	// 构造发送数据
	toAgentWeaponMsg := ujson.MsgWpReq{
		Id: args.Id,
		TaskId: args.TaskId,
		WpId: args.WeaponId,
		WpName: args.WeaponName,
		WpMd5: args.WeaponMd5,
		WpUrl: args.WeaponUrl,
		WpType: args.WeaponType,
		Timeout: args.Timeout,
		Params: args.Params,
	}
	content, _ := toAgentWeaponMsg.MarshalJSON()
	sendMsg := ujson.ProtoMsg{Type: ujson.ProtoWpMsgReq, Content: content}
	buf, err := sendMsg.MarshalJSON()
	if err != nil {
		return err
	}

	// 发送
	agent := client.AgentPoolInst.GetAgentByID(args.Id)
	err = agent.Connection.Write(buf)
	if err != nil {
		glog.Errorf("'%s' send '%s' error: %s", agent.IP, toAgentWeaponMsg.WpName, err.Error())
		return err
	}

	// 仅关注下发成功
	return nil
}

type ExecAction int

func init()  {
	rpc.Register(new(ExecAction))
}
